# encoding=utf-8
import torch
from thop import profile, clever_format
from ptflops import get_model_complexity_info
from pytorch_model_summary import summary
from networks import get_norm_layer, ResUNet

# load model
device = torch.device("cpu")
model = ResUNet(3, 3, 64, norm_layer=get_norm_layer())
model_parameter = torch.load("120_net_G_A.pth", map_location=device)
model.load_state_dict(model_parameter)
model.eval()

def use_thop():
    test_input = torch.randn(1, 3, 512, 512)
    flops, params = profile(model, inputs=(test_input, ))
    flops, params = clever_format([flops, params], '%.3f')
    print(f"运算量：{flops}, 参数量：{params}")

def use_ptflops():
    input_res = (3, 512, 512)
    macs, params = get_model_complexity_info(model, input_res, as_strings=True, print_per_layer_stat=True)
    print(f"模型 FLOPs: {macs}")
    print(f"模型参数量: {params}")

def use_summary():
    input_tensor = torch.rand(1, 3, 512, 512)
    model_summary = summary(model, input_tensor)
    print(model_summary)

if __name__ == '__main__':
    use_thop()
    # use_ptflops()
    # use_summary()